package store

import (
	"fmt"
	"log"
	"os"
	"path/filepath"

	"mig/pkg/collections/bloom_filter"
	"mig/pkg/types"
)

const (
	defaultFalsePositiveRate = 1e-5
	defaultMaxSize           = 10000
)

type JSONStore struct {
	workDir     string
	cache       map[string]*types.DigestTree
	bloomFilter *bloom_filter.BloomFilter
}

func NewJSONStore(workDir string) *JSONStore {
	m, k := bloom_filter.EnsureFalsePositiveProbabilityUnder(defaultFalsePositiveRate, defaultMaxSize)
	bf := bloom_filter.NewBloomFilter(m, k)
	// work the dir and construct the bloom filter
	// Create work directory if it doesn't exist
	if err := os.MkdirAll(workDir, 0755); err != nil {
		panic(fmt.Errorf("failed to create work directory: %v", err))
	}

	// Read all JSON files in the work directory and populate bloom filter
	files, err := os.ReadDir(workDir)
	if err != nil {
		panic(fmt.Errorf("failed to read work directory: %v", err))
	}

	for _, file := range files {
		if filepath.Ext(file.Name()) == ".json" {
			id := file.Name()[:len(file.Name())-5] // remove ".json" extension
			bf.Add([]byte(id))
		}
	}

	return &JSONStore{workDir: workDir, bloomFilter: bf, cache: make(map[string]*types.DigestTree)}
}

func (s *JSONStore) loadFromJSON(id string) (*types.DigestTree, error) {
	path := filepath.Join(s.workDir, fmt.Sprintf("%s.json", id))
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	dt, err := types.FromJSON(data)
	if err != nil {
		return nil, err
	}

	s.cache[id] = dt
	return dt, nil
}

func (s *JSONStore) GetDigestTree(id string) (*types.DigestTree, error) {
	if !s.bloomFilter.MightContain([]byte(id)) {
		return nil, nil
	}
	if dt, ok := s.cache[id]; ok {
		return dt, nil
	}

	return s.loadFromJSON(id)
}

func (s *JSONStore) StoreDigestTree(id string, dt *types.DigestTree) error {
	// Validate input
	if id == "" {
		return fmt.Errorf("id cannot be empty")
	}
	if dt == nil {
		log.Printf("digest tree cannot be nil")
		return nil
	}

	// Convert to JSON
	data := dt.ToJSON()

	// Create file path
	path := filepath.Join(s.workDir, fmt.Sprintf("%s.json", id))

	// Write to file with atomic write pattern
	tempPath := path + ".tmp"
	err := os.WriteFile(tempPath, []byte(data), 0644)
	if err != nil {
		return fmt.Errorf("failed to write temporary file: %w", err)
	}

	// Rename temp file to final file (atomic operation)
	err = os.Rename(tempPath, path)
	if err != nil {
		return fmt.Errorf("failed to rename temporary file: %w", err)
	}

	// Update cache
	s.cache[id] = dt

	// Update bloom filter
	s.bloomFilter.Add([]byte(id))

	return nil
}
